"""
脑波数据模型定义文件

本文件定义了脑波数据处理系统中模块间传递数据的标准格式。
主要包含时域数据和频域数据的结构定义，确保数据在不同模块间传递时的一致性和完整性。

使用方法：
- BrainWaveData: 主要数据容器，包含时域和频域数据
- TimeSeriesData: 512Hz时域数据结构
- FrequencyData: 1Hz频域数据结构（8个频段）
- DataValidator: 数据验证工具类

原理：
采用面向对象的数据结构设计，通过类型检查和数据验证确保数据质量。
支持可选的频域数据，满足不同处理阶段的需求。
"""

from typing import Optional, List, Dict, Any
from dataclasses import dataclass, field
from datetime import datetime
import numpy as np
from pydantic import ConfigDict


@dataclass
class TimeSeriesData:
    """
    时间序列数据模型
    用于存储原始EEG数据和相关元数据
    
    Attributes:
        data: numpy数组，形状为(n_samples,) 或 (n_channels, n_samples)
        sampling_rate: 采样率，默认512Hz
        timestamp: 数据采集时间戳
        channel_names: 通道名称列表，如['EEG']
        duration: 数据时长（秒）
    """
    data: np.ndarray
    sampling_rate: int = 512
    timestamp: datetime = field(default_factory=datetime.now)
    channel_names: List[str] = field(default_factory=lambda: ['EEG'])
    duration: float = 0.0
    
    def __post_init__(self):
        """初始化后处理，计算数据时长"""
        if self.data.ndim == 1:
            self.duration = len(self.data) / self.sampling_rate
        else:
            self.duration = self.data.shape[1] / self.sampling_rate


@dataclass
class FrequencyData:
    """
    频域数据结构 - 基于NeuroPy3的完整数据字段
    
    根据NeuroPy3代码分析，包含所有可获取的数据字段：
    - 8个频段功率数据 (ASIC_EEG_POWER)
    - 注意力和冥想指数
    - 信号质量和眨眼强度
    - 原始数据值
    
    Attributes:
        # 8个频段功率数据
        delta: Delta波段功率 (0.5-4Hz)
        theta: Theta波段功率 (4-8Hz) 
        low_alpha: 低Alpha波段功率 (8-10Hz)
        high_alpha: 高Alpha波段功率 (10-13Hz)
        low_beta: 低Beta波段功率 (13-17Hz)
        high_beta: 高Beta波段功率 (17-30Hz)
        low_gamma: 低Gamma波段功率 (30-40Hz)
        mid_gamma: 中Gamma波段功率 (40-100Hz)
        
        # NeuroSky算法输出
        attention: 注意力指数 (0-100)
        meditation: 冥想指数 (0-100)
        
        # 信号质量和其他指标
        poor_signal: 信号质量指标 (0=好信号, 200=无信号)
        blink_strength: 眨眼强度 (0-255)
        
        # 原始数据
        raw_value: 原始EEG数据值 (-32768 to 32767)
        
        # 元数据
        timestamp: 数据时间戳
        is_filtered: 是否为过滤后的数据
    """
    # 8个频段功率数据
    delta: float = 0.0
    theta: float = 0.0
    low_alpha: float = 0.0
    high_alpha: float = 0.0
    low_beta: float = 0.0
    high_beta: float = 0.0
    low_gamma: float = 0.0
    mid_gamma: float = 0.0
    
    # NeuroSky算法输出
    attention: float = 0.0
    meditation: float = 0.0
    
    # 信号质量和其他指标
    poor_signal: int = 0  # 0=好信号, 200=无信号
    blink_strength: int = 0  # 0-255
    
    # 原始数据
    raw_value: int = 0  # -32768 to 32767
    
    # 元数据
    timestamp: datetime = field(default_factory=datetime.now)
    is_filtered: bool = False  # 区分过滤前后的数据
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            # 8个频段功率数据
            'delta': self.delta,
            'theta': self.theta,
            'low_alpha': self.low_alpha,
            'high_alpha': self.high_alpha,
            'low_beta': self.low_beta,
            'high_beta': self.high_beta,
            'low_gamma': self.low_gamma,
            'mid_gamma': self.mid_gamma,
            
            # NeuroSky算法输出
            'attention': self.attention,
            'meditation': self.meditation,
            
            # 信号质量和其他指标
            'poor_signal': self.poor_signal,
            'blink_strength': self.blink_strength,
            
            # 原始数据
            'raw_value': self.raw_value,
            
            # 元数据
            'timestamp': self.timestamp.isoformat(),
            'is_filtered': self.is_filtered
        }


@dataclass
class BrainWaveData:
    """
    完整的脑波数据模型
    包含时间序列数据、频域数据和元数据
    
    Attributes:
        time_series: 时间序列数据
        frequency_data: 频域数据（可选）
        session_id: 会话ID
        room_id: 房间ID
        processing_stage: 处理阶段
        processing_timestamp: 处理时间戳
        metadata: 元数据字典
    """
    time_series: TimeSeriesData
    frequency_data: Optional[FrequencyData] = None
    session_id: str = ""
    room_id: str = ""
    processing_stage: str = "raw"  # raw, filtered, analyzed, mne_filtered, yasa_analyzed
    processing_timestamp: Optional[datetime] = None  # 处理时间戳
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def has_frequency_data(self) -> bool:
        """检查是否包含频域数据"""
        return self.frequency_data is not None
    
    def get_data_summary(self) -> Dict[str, Any]:
        """获取数据摘要信息"""
        summary = {
            'session_id': self.session_id,
            'processing_stage': self.processing_stage,
            'time_series_duration': self.time_series.duration,
            'time_series_samples': len(self.time_series.data) if self.time_series.data.ndim == 1 else self.time_series.data.shape[1],
            'sampling_rate': self.time_series.sampling_rate,
            'has_frequency_data': self.has_frequency_data(),
            'timestamp': self.time_series.timestamp.isoformat()
        }
        
        if self.has_frequency_data():
            summary['frequency_data'] = self.frequency_data.to_dict()
            
        return summary


class DataValidator:
    """
    数据验证工具类
    
    提供各种数据验证方法，确保数据格式的正确性和完整性。
    """
    
    @staticmethod
    def validate_time_series(data: TimeSeriesData) -> bool:
        """
        验证时域数据的有效性
        
        Args:
            data: 时域数据对象
            
        Returns:
            bool: 验证是否通过
        """
        try:
            # 检查数据类型
            if not isinstance(data.data, np.ndarray):
                return False
            
            # 检查数据维度
            if data.data.ndim not in [1, 2]:
                return False
            
            # 检查采样率
            if data.sampling_rate <= 0:
                return False
            
            # 检查数据长度
            if data.data.size == 0:
                return False
            
            # 检查通道名称数量
            if data.data.ndim == 2 and len(data.channel_names) != data.data.shape[0]:
                return False
            
            return True
            
        except Exception:
            return False
    
    @staticmethod
    def validate_frequency_data(data: FrequencyData) -> bool:
        """
        验证频域数据的有效性
        
        Args:
            data: 频域数据对象
            
        Returns:
            bool: 验证是否通过
        """
        try:
            # 检查所有频段数据是否为非负数
            frequency_values = [
                data.delta, data.theta, data.low_alpha, data.high_alpha,
                data.low_beta, data.high_beta, data.low_gamma, data.mid_gamma
            ]
            if any(val < 0 for val in frequency_values):
                return False
            
            # 检查注意力和冥想指数范围 (0-100)
            if not (0 <= data.attention <= 100) or not (0 <= data.meditation <= 100):
                return False
            
            # 检查信号质量范围 (0=好信号, 200=无信号)
            if not (0 <= data.poor_signal <= 200):
                return False
            
            # 检查眨眼强度范围 (0-255)
            if not (0 <= data.blink_strength <= 255):
                return False
            
            # 检查原始数据值范围 (-32768 to 32767)
            if not (-32768 <= data.raw_value <= 32767):
                return False
            
            return True
            
        except Exception:
            return False
    
    @staticmethod
    def validate_brain_wave_data(data: BrainWaveData) -> bool:
        """
        验证完整脑波数据的有效性
        
        Args:
            data: 脑波数据对象
            
        Returns:
            bool: 验证是否通过
        """
        try:
            # 验证时域数据（必需）
            if not DataValidator.validate_time_series(data.time_series):
                return False
            
            # 验证频域数据（如果存在）
            if data.frequency_data is not None:
                if not DataValidator.validate_frequency_data(data.frequency_data):
                    return False
            
            # 验证处理阶段标识
            valid_stages = ['raw', 'filtered', 'analyzed']
            if data.processing_stage not in valid_stages:
                return False
            
            return True
            
        except Exception:
            return False


# 数据格式常量定义
class DataConstants:
    """数据格式相关常量"""
    
    # 采样率
    TIME_SERIES_SAMPLING_RATE = 512  # Hz
    FREQUENCY_UPDATE_RATE = 1  # Hz
    
    # 频段定义 - 基于NeuroPy3的8个频段
    FREQUENCY_BANDS = {
        'delta': (0.5, 4.0),
        'theta': (4.0, 8.0),
        'low_alpha': (8.0, 10.0),
        'high_alpha': (10.0, 13.0),
        'low_beta': (13.0, 17.0),
        'high_beta': (17.0, 30.0),
        'low_gamma': (30.0, 40.0),
        'mid_gamma': (40.0, 100.0)
    }
    
    # 处理阶段
    PROCESSING_STAGES = ['raw', 'filtered', 'analyzed']
    
    # 数据质量等级 - 基于poor_signal值
    SIGNAL_EXCELLENT = 0    # 优秀信号
    SIGNAL_GOOD = 25       # 良好信号  
    SIGNAL_FAIR = 50       # 一般信号
    SIGNAL_POOR = 100      # 差信号
    SIGNAL_NO_CONTACT = 200 # 无接触